package com.xypower.graph;

import java.util.*;
//图转二叉树，最小生成树,P算法
public class Prim {
    public static class EdgeComparator implements Comparator<Edge>{

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    public static Set<Edge> primMST(Graph graph){
        //解锁的边进入小根堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

        HashSet<Node> set = new HashSet<>();
        Set<Edge> result = new HashSet<>();//依次挑选的边在result里

        for(Node node : graph.nodes.values()){//随便挑一个点
            if (!set.contains(node)){
                set.add(node);
                for (Edge edge : node.edges){//由一个点，解锁所有相连的边
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()){
                    Edge edge = priorityQueue.poll();//弹出解锁的边中最小的边
                    Node toNode = edge.to;              //可能的一个新的点
                    if (!set.contains(toNode)){         //不包含的时候就是新的点
                        set.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges){
                            priorityQueue.add(nextEdge);
                        }
                    }
                }
            }
          //break;
        }
        return result;
    }
}
